import tensorrt as trt
import common

# Or using trtexec:
# trtexec --verbose --onnx=resnet101_base.onnx --explicitBatch --saveEngine=resnet101_base_fp32_3.trt \
# --minShapes=src:1x3x224x224,bgr:1x3x224x224 \
# --optShapes=src:1x3x1080x1920,bgr:1x3x1080x1920 \
# --maxShapes=src:1x3x1080x1920,bgr:1x3x1080x1920 \
# --fp16

import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--batch-size', type=int, default=3)
parser.add_argument('--mode', type=str, default='fp16', choices=['fp32', 'fp16'])
args = parser.parse_args()

batch_size = args.batch_size
mode = args.mode
print("batch size:", batch_size, "mode:", mode)

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

class ModelData(object):
    MODEL_FILE = "./models/resnet101_base.onnx"


def build_engine(model_file, mode='fp32', int8_calibrator=None):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(network,
                                                                                                               TRT_LOGGER) as parser:
        assert mode.lower() in ['fp32', 'fp16', 'int8'], "mode should be in ['fp32', 'fp16', 'int8']"
        builder.max_workspace_size = common.GiB(15)
        builder.max_batch_size = batch_size

        # Parse the onnx network
        with open(model_file, 'rb') as model:
            parser.parse(model.read())
        config = builder.create_builder_config()
        profile = builder.create_optimization_profile()
        profile.set_shape('src', (1, 3, 224, 224), (batch_size, 3, 1080, 1920), (batch_size, 3, 1080, 1920))
        profile.set_shape('bgr', (1, 3, 224, 224), (batch_size, 3, 1080, 1920), (batch_size, 3, 1080, 1920))
        config.add_optimization_profile(profile)
        if mode.lower() == 'int8':
            # assert (builder.platform_has_fast_int8 == True), "not support int8"
            # builder.int8_mode = True
            # builder.int8_calibrator = int8_calibrator
            pass
        elif mode.lower() == 'fp16':
            assert (builder.platform_has_fast_fp16 == True), "not support fp16"
            config.flags = 1 << int(trt.BuilderFlag.FP16)
        return builder.build_engine(network, config)


def gen_engine():
    with build_engine(ModelData.MODEL_FILE, mode=mode) as engine:
        # Build an engine, allocate buffers and create a stream.
        with open('./models/resnet101_base_' + mode + '_' + str(batch_size) + '.trt', 'wb') as f:
            print("Start serializing...")
            f.write(engine.serialize())
            print("Engine file is Saved.")
        return

gen_engine()
